{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4728ae6f",
   "metadata": {},
   "source": [
    "# mm-relu 示例"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44202d8",
   "metadata": {},
   "source": [
    "记\n",
    "\n",
    "$$\n",
    "\\begin{cases}\n",
    "\\mathbf{x}_i = \\begin{bmatrix} \\mathbf{x}_i^{(1)} & \\mathbf{x}_i^{(2)} & \\cdots & \\mathbf{x}_i^{(m)} \\end{bmatrix}^T  \\\\\n",
    "\\mathbf{y}_j = \\begin{bmatrix} \\mathbf{y}_j^{(1)} & \\mathbf{y}_j^{(2)} & \\cdots & \\mathbf{y}_j^{(p)} \\end{bmatrix}^T \n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "有\n",
    "\n",
    "\n",
    "$$\n",
    "\\begin{cases}\n",
    "<\\mathbf{x}_i, \\mathbf{y}_j> = \\mathbf{x}_i^T \\mathbf{y}_j \\\\\n",
    "\\mathbf{X} = \\begin{bmatrix} \\mathbf{x}_1^T & \\mathbf{x}_2^T & \\cdots & \\mathbf{x}_n^T \\end{bmatrix}^T \\\\\n",
    "\\mathbf{Y} = \\begin{bmatrix} \\mathbf{y}_1^T & \\mathbf{y}_2^T & \\cdots & \\mathbf{y}_n^T \\end{bmatrix}^T\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "可以推出\n",
    "\n",
    "$$\n",
    "<\\mathbf{X}, \\mathbf{Y}> = \\mathbf{X}^T \\mathbf{Y}  = (\\mathbf{x}_i^T \\mathbf{y}_j)_{m \\times p}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3077a4a6",
   "metadata": {},
   "source": [
    "ReLU 函数定义：\n",
    "\n",
    "$$\n",
    "\\mathbf{relu}(X) = \\mathbf{\\max}(X, 0)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a35b9625",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8511a417",
   "metadata": {},
   "source": [
    "使用 NumPy 实现如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2073fdb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = \"float32\"\n",
    "a_np = np.random.rand(128, 128).astype(dtype)\n",
    "b_np = np.random.rand(128, 128).astype(dtype)\n",
    "# a @ b 等价于 np.matmul(a, b)\n",
    "c_mm_relu = np.maximum(a_np @ b_np, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66cef2f0",
   "metadata": {},
   "source": [
    "NumPy 低级实现："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6e86c5eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):\n",
    "    # 分配中间数组用于存储矩阵乘法的结果\n",
    "    Y = np.empty((128, 128), dtype=\"float32\")\n",
    "    for i in range(128):\n",
    "        for j in range(128):\n",
    "            for k in range(128):\n",
    "                if k == 0:\n",
    "                    Y[i, j] = 0\n",
    "                Y[i, j] = Y[i, j] + A[i, k] * B[k, j]\n",
    "    for i in range(128):\n",
    "        for j in range(128):\n",
    "            C[i, j] = max(Y[i, j], 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0ba2208",
   "metadata": {},
   "source": [
    "验证数值一致性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e2961164",
   "metadata": {},
   "outputs": [],
   "source": [
    "c_np = np.empty((128, 128), dtype=dtype)\n",
    "lnumpy_mm_relu(a_np, b_np, c_np)\n",
    "np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5f8496c",
   "metadata": {},
   "source": [
    "## TensorIR 实现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "08af76dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModule:\n",
    "    @T.prim_func\n",
    "    def mm_relu(A: T.Buffer((128, 128), \"float32\"),\n",
    "                B: T.Buffer((128, 128), \"float32\"),\n",
    "                C: T.Buffer((128, 128), \"float32\")):\n",
    "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
    "        Y = T.alloc_buffer((128, 128), dtype=\"float32\")\n",
    "        for i, j, k in T.grid(128, 128, 128):\n",
    "            with T.block(\"Y\"):\n",
    "                vi = T.axis.spatial(128, i)\n",
    "                vj = T.axis.spatial(128, j)\n",
    "                vk = T.axis.reduce(128, k)\n",
    "                with T.init():\n",
    "                    Y[vi, vj] = T.float32(0)\n",
    "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
    "        for i, j in T.grid(128, 128):\n",
    "            with T.block(\"C\"):\n",
    "                vi = T.axis.spatial(128, i)\n",
    "                vj = T.axis.spatial(128, j)\n",
    "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e209fae0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tir.BlockRV(0x55bcd40cc720)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sch = tvm.tir.Schedule(MyModule)\n",
    "block_Y = sch.get_block(\"Y\", func_name=\"mm_relu\")\n",
    "i, j, k = sch.get_loops(block_Y)\n",
    "j0, j1 = sch.split(j, factors=[None, 4])\n",
    "sch.reorder(j0, k, j1)\n",
    "block_C = sch.get_block(\"C\", \"mm_relu\")\n",
    "sch.reverse_compute_at(block_C, j0)\n",
    "sch.decompose_reduction(block_Y, k) # 将 Y 元素的初始化与归约更新分开"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e620986f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">mm_relu</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        Y <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">32</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> j_1_init <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">4</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;Y_init&quot;</span>):\n",
       "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
       "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1_init)\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads()\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
       "                    Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> k, j_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">4</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;Y_update&quot;</span>):\n",
       "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
       "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1)\n",
       "                    vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">128</span>, k)\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Y[vi, vj], A[vi, vk], B[vk, vj])\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
       "                    Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[vk, vj]\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> ax0 <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">4</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
       "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> ax0)\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Y[vi, vj])\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi, vj])\n",
       "                    C[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(Y[vi, vj], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e3f9fe96",
   "metadata": {},
   "outputs": [],
   "source": [
    "rt_lib = tvm.build(MyModule, target=\"llvm\")\n",
    "a_nd = tvm.nd.array(a_np)\n",
    "b_nd = tvm.nd.array(b_np)\n",
    "c_nd = tvm.nd.empty((128, 128), dtype=\"float32\")\n",
    "func_mm_relu = rt_lib[\"mm_relu\"]\n",
    "func_mm_relu(a_nd, b_nd, c_nd)\n",
    "\n",
    "np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "af92d3b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# from tvm import tir\n",
      "def apply_trace(sch: tir.Schedule) -> None:\n",
      "  b0 = sch.get_block(name=\"Y\", func_name=\"mm_relu\")\n",
      "  l1, l2, l3 = sch.get_loops(block=b0)\n",
      "  l4, l5 = sch.split(loop=l2, factors=[None, 4], preserve_unit_iters=True, disable_predication=False)\n",
      "  sch.reorder(l4, l3, l5)\n",
      "  b6 = sch.get_block(name=\"C\", func_name=\"mm_relu\")\n",
      "  sch.reverse_compute_at(block=b6, loop=l4, preserve_unit_loops=False, index=-1)\n",
      "  b7 = sch.decompose_reduction(block=b0, loop=l3)\n"
     ]
    }
   ],
   "source": [
    "print(sch.trace)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bba5797a",
   "metadata": {},
   "source": [
    "## 随机调度变换 (Stochastic Schedule Transformation)\n",
    "\n",
    "考虑简单的模型："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e813dc5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModule:\n",
    "    @T.prim_func\n",
    "    def main(\n",
    "        A: T.Buffer((128, 128), \"float32\"),\n",
    "        B: T.Buffer((128, 128), \"float32\"),\n",
    "        C: T.Buffer((128, 128), \"float32\"),\n",
    "    ):\n",
    "        T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
    "        for i, j, k in T.grid(128, 128, 128):\n",
    "            with T.block(\"C\"):\n",
    "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
    "                with T.init():\n",
    "                    C[vi, vj] = 0.0\n",
    "                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3b0c61b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">64</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> j_1_init <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">2</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C_init&quot;</span>):\n",
       "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
       "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1_init)\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads()\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi, vj])\n",
       "                    C[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> k, j_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">2</span>):\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C_update&quot;</span>):\n",
       "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
       "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1)\n",
       "                    vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">128</span>, k)\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(C[vi, vj], A[vi, vk], B[vk, vj])\n",
       "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi, vj])\n",
       "                    C[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[vk, vj]\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def stochastic_schedule_mm(sch: tvm.tir.Schedule):\n",
    "    block_C = sch.get_block(\"C\", \"main\")\n",
    "    i, j, k = sch.get_loops(block=block_C)\n",
    "    j_factors = sch.sample_perfect_tile(loop=j, n=2)\n",
    "    j_0, j_1 = sch.split(loop=j, factors=j_factors)\n",
    "    sch.reorder(i, j_0, k, j_1)\n",
    "    sch.decompose_reduction(block_C, k)\n",
    "    return sch\n",
    "\n",
    "sch = tvm.tir.Schedule(MyModule)\n",
    "sch = stochastic_schedule_mm(sch)\n",
    "sch.mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ad4482b",
   "metadata": {},
   "source": [
    "## 随机变换搜索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "00a09a16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Attempt 0, time-cost: 0.379 ms====\n",
      "# from tvm import tir\n",
      "def apply_trace(sch: tir.Schedule) -> None:\n",
      "  b0 = sch.get_block(name=\"C\", func_name=\"main\")\n",
      "  l1, l2, l3 = sch.get_loops(block=b0)\n",
      "  v4, v5 = sch.sample_perfect_tile(loop=l2, n=2, max_innermost_factor=16, decision=[8, 16])\n",
      "  l6, l7 = sch.split(loop=l2, factors=[v4, v5], preserve_unit_iters=True, disable_predication=False)\n",
      "  sch.reorder(l1, l6, l3, l7)\n",
      "  b8 = sch.decompose_reduction(block=b0, loop=l3)\n",
      "=====Attempt 1, time-cost: 0.363 ms====\n",
      "# from tvm import tir\n",
      "def apply_trace(sch: tir.Schedule) -> None:\n",
      "  b0 = sch.get_block(name=\"C\", func_name=\"main\")\n",
      "  l1, l2, l3 = sch.get_loops(block=b0)\n",
      "  v4, v5 = sch.sample_perfect_tile(loop=l2, n=2, max_innermost_factor=16, decision=[16, 8])\n",
      "  l6, l7 = sch.split(loop=l2, factors=[v4, v5], preserve_unit_iters=True, disable_predication=False)\n",
      "  sch.reorder(l1, l6, l3, l7)\n",
      "  b8 = sch.decompose_reduction(block=b0, loop=l3)\n",
      "=====Attempt 2, time-cost: 2.229 ms====\n",
      "# from tvm import tir\n",
      "def apply_trace(sch: tir.Schedule) -> None:\n",
      "  b0 = sch.get_block(name=\"C\", func_name=\"main\")\n",
      "  l1, l2, l3 = sch.get_loops(block=b0)\n",
      "  v4, v5 = sch.sample_perfect_tile(loop=l2, n=2, max_innermost_factor=16, decision=[128, 1])\n",
      "  l6, l7 = sch.split(loop=l2, factors=[v4, v5], preserve_unit_iters=True, disable_predication=False)\n",
      "  sch.reorder(l1, l6, l3, l7)\n",
      "  b8 = sch.decompose_reduction(block=b0, loop=l3)\n",
      "=====Attempt 3, time-cost: 1.015 ms====\n",
      "# from tvm import tir\n",
      "def apply_trace(sch: tir.Schedule) -> None:\n",
      "  b0 = sch.get_block(name=\"C\", func_name=\"main\")\n",
      "  l1, l2, l3 = sch.get_loops(block=b0)\n",
      "  v4, v5 = sch.sample_perfect_tile(loop=l2, n=2, max_innermost_factor=16, decision=[64, 2])\n",
      "  l6, l7 = sch.split(loop=l2, factors=[v4, v5], preserve_unit_iters=True, disable_predication=False)\n",
      "  sch.reorder(l1, l6, l3, l7)\n",
      "  b8 = sch.decompose_reduction(block=b0, loop=l3)\n",
      "=====Attempt 4, time-cost: 1.853 ms====\n",
      "# from tvm import tir\n",
      "def apply_trace(sch: tir.Schedule) -> None:\n",
      "  b0 = sch.get_block(name=\"C\", func_name=\"main\")\n",
      "  l1, l2, l3 = sch.get_loops(block=b0)\n",
      "  v4, v5 = sch.sample_perfect_tile(loop=l2, n=2, max_innermost_factor=16, decision=[128, 1])\n",
      "  l6, l7 = sch.split(loop=l2, factors=[v4, v5], preserve_unit_iters=True, disable_predication=False)\n",
      "  sch.reorder(l1, l6, l3, l7)\n",
      "  b8 = sch.decompose_reduction(block=b0, loop=l3)\n"
     ]
    }
   ],
   "source": [
    "def random_search(mod: tvm.IRModule, num_trials=5):\n",
    "    best_result = None\n",
    "    best_sch = None\n",
    "\n",
    "    for i in range(num_trials):\n",
    "        sch = stochastic_schedule_mm(tvm.tir.Schedule(mod))\n",
    "        lib = tvm.build(sch.mod, target=\"llvm\")\n",
    "        f_timer_after = lib.time_evaluator(\"main\", tvm.cpu())\n",
    "        result = f_timer_after(a_nd, b_nd, c_nd).mean\n",
    "\n",
    "        print(\"=====Attempt %d, time-cost: %.3f ms====\" % (i, result * 1000))\n",
    "        print(sch.trace)\n",
    "\n",
    "        # book keep the best result so far\n",
    "        if best_result is None or result < best_result:\n",
    "            best_result = result\n",
    "            best_sch = sch\n",
    "\n",
    "    return best_sch\n",
    "\n",
    "sch = random_search(MyModule)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29ba0050",
   "metadata": {},
   "source": [
    "使用随机变换来指定好的程序的搜索空间，使用 ``tune_tir`` API 帮助在搜索空间内搜索并找到最优的调度变换。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f412afa2",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-04-10 13:50:25 [INFO] Logging directory: .temp/tune_tmp/logs\n",
      "2025-04-10 13:50:42 [INFO] LocalBuilder: max_workers = 24\n",
      "2025-04-10 13:50:42 [INFO] LocalRunner: max_workers = 1\n",
      "2025-04-10 13:50:43 [INFO] [task_scheduler.cc:159] Initializing Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1               N/A             N/A    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                     N/A         0           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n",
      "2025-04-10 13:50:53 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n",
      "2025-04-10 13:50:53 [INFO] [task_scheduler.cc:180] TaskScheduler picks Task #0: \"main\"\n",
      "2025-04-10 13:50:54 [INFO] [task_scheduler.cc:193] Sending 5 sample(s) to builder\n",
      "2025-04-10 13:50:55 [INFO] [task_scheduler.cc:195] Sending 5 sample(s) to runner\n",
      "2025-04-10 13:51:22 [DEBUG] XGB iter   0: tr-p-rmse: 0.327315\ttr-a-peak@32: 1.000000\ttr-rmse: 0.356233\ttr-rmse: 0.356233\n",
      "2025-04-10 13:51:22 [DEBUG] XGB iter  25: tr-p-rmse: 0.106004\ttr-a-peak@32: 1.000000\ttr-rmse: 0.062224\ttr-rmse: 0.062224\n",
      "2025-04-10 13:51:22 [DEBUG] XGB iter  50: tr-p-rmse: 0.101582\ttr-a-peak@32: 1.000000\ttr-rmse: 0.059773\ttr-rmse: 0.059773\n",
      "2025-04-10 13:51:22 [DEBUG] XGB iter  75: tr-p-rmse: 0.101558\ttr-a-peak@32: 1.000000\ttr-rmse: 0.059773\ttr-rmse: 0.059773\n",
      "2025-04-10 13:51:22 [DEBUG] XGB iter 100: tr-p-rmse: 0.101558\ttr-a-peak@32: 1.000000\ttr-rmse: 0.059773\ttr-rmse: 0.059773\n",
      "2025-04-10 13:51:22 [DEBUG] XGB stopped. Best iteration: [74] tr-p-rmse:0.10156\ttr-a-peak@32:1.00000\ttr-rmse:0.05977\ttr-rmse:0.05977 \n",
      "2025-04-10 13:51:22 [INFO] [task_scheduler.cc:237] [Updated] Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>13.2337</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>5</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1           13.2337        316.9420    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                316.9420         5           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-04-10 13:51:22 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |        13.2337 |     316.9420 |              316.9420 |      5 |      \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "2025-04-10 13:51:22 [INFO] [task_scheduler.cc:180] TaskScheduler picks Task #0: \"main\"\n",
      "2025-04-10 13:51:23 [INFO] [task_scheduler.cc:193] Sending 0 sample(s) to builder\n",
      "2025-04-10 13:51:23 [INFO] [task_scheduler.cc:195] Sending 0 sample(s) to runner\n",
      "2025-04-10 13:51:23 [INFO] [task_scheduler.cc:237] [Updated] Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>13.2337</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>5</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1           13.2337        316.9420    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                316.9420         5           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-04-10 13:51:23 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |        13.2337 |     316.9420 |              316.9420 |      5 |      \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "2025-04-10 13:51:23 [INFO] [task_scheduler.cc:180] TaskScheduler picks Task #0: \"main\"\n",
      "2025-04-10 13:51:23 [INFO] [task_scheduler.cc:193] Sending 0 sample(s) to builder\n",
      "2025-04-10 13:51:23 [INFO] [task_scheduler.cc:195] Sending 0 sample(s) to runner\n",
      "2025-04-10 13:51:23 [INFO] [task_scheduler.cc:237] [Updated] Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>13.2337</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>5</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1           13.2337        316.9420    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                316.9420         5           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-04-10 13:51:23 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |        13.2337 |     316.9420 |              316.9420 |      5 |      \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "2025-04-10 13:51:23 [INFO] [task_scheduler.cc:180] TaskScheduler picks Task #0: \"main\"\n",
      "2025-04-10 13:51:24 [INFO] [task_scheduler.cc:193] Sending 0 sample(s) to builder\n",
      "2025-04-10 13:51:24 [INFO] [task_scheduler.cc:195] Sending 0 sample(s) to runner\n",
      "2025-04-10 13:51:24 [INFO] [task_scheduler.cc:237] [Updated] Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>13.2337</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>5</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1           13.2337        316.9420    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                316.9420         5           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "2025-04-10 13:51:24 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |        13.2337 |     316.9420 |              316.9420 |      5 |      \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "2025-04-10 13:51:24 [INFO] [task_scheduler.cc:180] TaskScheduler picks Task #0: \"main\"\n",
      "2025-04-10 13:51:25 [INFO] [task_scheduler.cc:193] Sending 0 sample(s) to builder\n",
      "2025-04-10 13:51:25 [INFO] [task_scheduler.cc:195] Sending 0 sample(s) to runner\n",
      "2025-04-10 13:51:25 [INFO] [task_scheduler.cc:237] [Updated] Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>13.2337</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>5</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1           13.2337        316.9420    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                316.9420         5           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-04-10 13:51:25 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |        13.2337 |     316.9420 |              316.9420 |      5 |      \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "2025-04-10 13:51:25 [INFO] [task_scheduler.cc:180] TaskScheduler picks Task #0: \"main\"\n",
      "2025-04-10 13:51:25 [INFO] [task_scheduler.cc:260] Task #0 has finished. Remaining task(s): 0\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>13.2337</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>316.9420</td>\n",
       "      <td>5</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1           13.2337        316.9420    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                316.9420         5       Y   "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-04-10 13:51:25 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |        13.2337 |     316.9420 |              316.9420 |      5 |    Y \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n",
      "\n",
      "Total trials: 5\n",
      "Total latency (us): 316.942\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from tvm import meta_schedule as ms\n",
    "\n",
    "database = ms.tune_tir(\n",
    "    mod=MyModule,\n",
    "    target=\"llvm --num-cores=1\",\n",
    "    max_trials_global=64,\n",
    "    num_trials_per_iter=64,\n",
    "    space=ms.space_generator.ScheduleFn(stochastic_schedule_mm),\n",
    "    work_dir=\".temp/tune_tmp\",\n",
    ")\n",
    "\n",
    "sch = ms.tir_integration.compile_tir(database, MyModule, \"llvm --num-cores=1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "850cdd2d",
   "metadata": {},
   "source": [
    "查看调优结果:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6c3ca9f5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm import tir</span>\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">apply_trace</span>(sch: tir<span style=\"color: #AA22FF; font-weight: bold\">.</span>Schedule) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
       "  b0 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_block(name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;C&quot;</span>, func_name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;main&quot;</span>)\n",
       "  l1, l2, l3 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_loops(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b0)\n",
       "  v4, v5 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>sample_perfect_tile(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l2, n<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">2</span>, max_innermost_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">16</span>, decision<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>])\n",
       "  l6, l7 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>split(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l2, factors<span style=\"color: #AA22FF; font-weight: bold\">=</span>[v4, v5], preserve_unit_iters<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>, disable_predication<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>reorder(l1, l6, l3, l7)\n",
       "  b8 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>decompose_reduction(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b0, loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l3)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>enter_postproc()\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.trace.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "11c0fc4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time cost of MyModule after tuning: 0.372 ms\n"
     ]
    }
   ],
   "source": [
    "lib = tvm.build(sch.mod, target=\"llvm\")\n",
    "f_timer_after = lib.time_evaluator(\"main\", tvm.cpu())\n",
    "print(\"Time cost of MyModule after tuning: %.3f ms\" % (f_timer_after(a_nd, b_nd, c_nd).mean * 1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ad1a544",
   "metadata": {},
   "source": [
    "## 利用默认的自动调度\n",
    "\n",
    "Meta-Schedule 带有内置通用随机变换集合，能够适用于广泛的 TensorIR 计算。这种方法也称为自动调度 (auto-scheduling)，因为搜索空间是由系统生成的。可以通过删除行 `space=ms.space_generator.ScheduleFn(stochastic_schedule_mm)` 来运行它。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94bff82f",
   "metadata": {},
   "source": [
    "在底层，Meta-Schedule 分析每个 TensorIR block 的数据访问和循环模式，并提出对程序的随机变换方式。我们不会在本章中讨论这些通用的变换，但要注意它们也只是随机转换加上代码分析而已。可以使用上一节中学到的相同机制来增强自动调度。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "63c67d1f",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-04-10 14:11:05 [INFO] Logging directory: .temp/tune_tmp/logs\n",
      "2025-04-10 14:11:05 [INFO] LocalBuilder: max_workers = 24\n",
      "2025-04-10 14:11:06 [INFO] LocalRunner: max_workers = 1\n",
      "2025-04-10 14:11:07 [INFO] [task_scheduler.cc:159] Initializing Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1               N/A             N/A    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                     N/A         0           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n",
      "2025-04-10 14:11:07 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n",
      "2025-04-10 14:11:07 [INFO] [task_scheduler.cc:180] TaskScheduler picks Task #0: \"main\"\n",
      "2025-04-10 14:11:08 [INFO] [task_scheduler.cc:193] Sending 64 sample(s) to builder\n",
      "2025-04-10 14:11:15 [INFO] [task_scheduler.cc:195] Sending 64 sample(s) to runner\n",
      "2025-04-10 14:11:30 [DEBUG] XGB iter   0: tr-p-rmse: 0.428030\ttr-a-peak@32: 0.998691\ttr-rmse: 0.277890\ttr-rmse: 0.277890\n",
      "2025-04-10 14:11:30 [DEBUG] XGB iter  25: tr-p-rmse: 0.049583\ttr-a-peak@32: 1.000000\ttr-rmse: 0.327711\ttr-rmse: 0.327711\n",
      "2025-04-10 14:11:30 [DEBUG] XGB iter  50: tr-p-rmse: 0.049571\ttr-a-peak@32: 1.000000\ttr-rmse: 0.327729\ttr-rmse: 0.327729\n",
      "2025-04-10 14:11:30 [DEBUG] XGB iter  75: tr-p-rmse: 0.049571\ttr-a-peak@32: 1.000000\ttr-rmse: 0.327729\ttr-rmse: 0.327729\n",
      "2025-04-10 14:11:30 [DEBUG] XGB stopped. Best iteration: [33] tr-p-rmse:0.04957\ttr-a-peak@32:1.00000\ttr-rmse:0.32773\ttr-rmse:0.32773 \n",
      "2025-04-10 14:11:30 [INFO] [task_scheduler.cc:237] [Updated] Task #0: \"main\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>184.1318</td>\n",
       "      <td>22.7788</td>\n",
       "      <td>22.7788</td>\n",
       "      <td>64</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1          184.1318         22.7788    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                 22.7788        64           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total trials: 64\n",
      "Total latency (us): 22.7788\n",
      "\n",
      "2025-04-10 14:11:30 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |       184.1318 |      22.7788 |               22.7788 |     64 |      \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 64\n",
      "Total latency (us): 22.7788\n",
      "\n",
      "2025-04-10 14:11:30 [INFO] [task_scheduler.cc:260] Task #0 has finished. Remaining task(s): 0\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>main</td>\n",
       "      <td>4194304</td>\n",
       "      <td>1</td>\n",
       "      <td>184.1318</td>\n",
       "      <td>22.7788</td>\n",
       "      <td>22.7788</td>\n",
       "      <td>64</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Name       FLOP    Weight    Speed (GFLOPS)    Latency (us)   \\\n",
       "0   main    4194304         1          184.1318         22.7788    \n",
       "\n",
       "    Weighted Latency (us)    Trials    Done   \n",
       "0                 22.7788        64       Y   "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total trials: 64\n",
      "Total latency (us): 22.7788\n",
      "\n",
      "2025-04-10 14:11:30 [DEBUG] [task_scheduler.cc:318] \n",
      " ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "------------------------------------------------------------------------------------------------------\n",
      "  0 | main | 4194304 |      1 |       184.1318 |      22.7788 |               22.7788 |     64 |    Y \n",
      "------------------------------------------------------------------------------------------------------\n",
      "Total trials: 64\n",
      "Total latency (us): 22.7788\n",
      "\n"
     ]
    }
   ],
   "source": [
    "database = ms.tune_tir(\n",
    "    mod=MyModule,\n",
    "    target=\"llvm --num-cores=1\",\n",
    "    max_trials_global=64,\n",
    "    num_trials_per_iter=64,\n",
    "    work_dir=\".temp/tune_tmp\",\n",
    ")\n",
    "sch = ms.tir_integration.compile_tir(database, MyModule, \"llvm --num-cores=1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "88bacec2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time cost of MyModule after tuning: 0.040 ms\n"
     ]
    }
   ],
   "source": [
    "lib = tvm.build(sch.mod, target=\"llvm\")\n",
    "f_timer_after = lib.time_evaluator(\"main\", tvm.cpu())\n",
    "print(\"Time cost of MyModule after tuning: %.3f ms\" % (f_timer_after(a_nd, b_nd, c_nd).mean * 1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cd9ee8f",
   "metadata": {},
   "source": [
    "结果比原始代码快得多。可以查看历史轨迹和最终代码。在高层次的理解中，历史轨迹包含：\n",
    "- 更多级的循环变换\n",
    "- 中间计算的矢量化\n",
    "- 并行化和循环展开"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "5ab0dd0e",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm import tir</span>\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">apply_trace</span>(sch: tir<span style=\"color: #AA22FF; font-weight: bold\">.</span>Schedule) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
       "  b0 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_block(name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;C&quot;</span>, func_name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;main&quot;</span>)\n",
       "  b1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_block(name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;root&quot;</span>, func_name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;main&quot;</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>annotate(block_or_loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>b0, ann_key<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;meta_schedule.tiling_structure&quot;</span>, ann_val<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;SSRSRS&quot;</span>)\n",
       "  l2, l3, l4 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_loops(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b0)\n",
       "  v5, v6, v7, v8 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>sample_perfect_tile(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l2, n<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">4</span>, max_innermost_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">64</span>, decision<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>])\n",
       "  l9, l10, l11, l12 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>split(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l2, factors<span style=\"color: #AA22FF; font-weight: bold\">=</span>[v5, v6, v7, v8], preserve_unit_iters<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>, disable_predication<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>)\n",
       "  v13, v14, v15, v16 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>sample_perfect_tile(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l3, n<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">4</span>, max_innermost_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">64</span>, decision<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>])\n",
       "  l17, l18, l19, l20 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>split(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l3, factors<span style=\"color: #AA22FF; font-weight: bold\">=</span>[v13, v14, v15, v16], preserve_unit_iters<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>, disable_predication<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>)\n",
       "  v21, v22 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>sample_perfect_tile(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l4, n<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">2</span>, max_innermost_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">64</span>, decision<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">2</span>])\n",
       "  l23, l24 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>split(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l4, factors<span style=\"color: #AA22FF; font-weight: bold\">=</span>[v21, v22], preserve_unit_iters<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>, disable_predication<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>reorder(l9, l17, l10, l18, l23, l11, l19, l24, l12, l20)\n",
       "  b25 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>cache_write(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b0, write_buffer_index<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>, storage_scope<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>reverse_compute_at(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b25, loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l18, preserve_unit_loops<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>, index<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">1</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>annotate(block_or_loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>b1, ann_key<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;meta_schedule.parallel&quot;</span>, ann_val<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">16</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>annotate(block_or_loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>b1, ann_key<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;meta_schedule.vectorize&quot;</span>, ann_val<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">64</span>)\n",
       "  v26 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>sample_categorical(candidates<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">512</span>], probs<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0.25</span>, <span style=\"color: #008000\">0.25</span>, <span style=\"color: #008000\">0.25</span>, <span style=\"color: #008000\">0.25</span>], decision<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">2</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>annotate(block_or_loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>b1, ann_key<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;meta_schedule.unroll_explicit&quot;</span>, ann_val<span style=\"color: #AA22FF; font-weight: bold\">=</span>v26)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>enter_postproc()\n",
       "  b27 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_block(name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;root&quot;</span>, func_name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;main&quot;</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>unannotate(block_or_loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>b27, ann_key<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;meta_schedule.parallel&quot;</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>unannotate(block_or_loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>b27, ann_key<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;meta_schedule.vectorize&quot;</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>unannotate(block_or_loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>b27, ann_key<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;meta_schedule.unroll_explicit&quot;</span>)\n",
       "  b28, b29 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_child_blocks(b27)\n",
       "  l30, l31, l32, l33, l34, l35, l36, l37, l38, l39 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_loops(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b28)\n",
       "  l40 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>fuse(l30, l31, l32, preserve_unit_iters<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>parallel(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l40)\n",
       "  l41 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>fuse(l39, preserve_unit_iters<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>vectorize(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l41)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>annotate(block_or_loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l40, ann_key<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;pragma_auto_unroll_max_step&quot;</span>, ann_val<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">64</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>annotate(block_or_loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l40, ann_key<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;pragma_unroll_explicit&quot;</span>, ann_val<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>)\n",
       "  l42, l43, l44, l45 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_loops(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b29)\n",
       "  l46 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>fuse(l45, preserve_unit_iters<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>vectorize(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l46)\n",
       "  b47 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_block(name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;C&quot;</span>, func_name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;main&quot;</span>)\n",
       "  l48, l49, l50, l51, l52, l53, l54, l55 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_loops(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b47)\n",
       "  b56 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>decompose_reduction(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b47, loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l50)\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sch.trace.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dacbeac3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
